# -*- coding: utf-8 -*-
"""
PBAR_CompareZspecRadiographyTransmission.py

Open basic scan radiography and zspec images

Created on Mon May 05 11:17:35 2014

@author: jkwong
"""

import copy, glob, shutil, codecs, os
import PBAR_Zspec, PBAR_Cargo
reload(PBAR_Zspec)
reload(PBAR_Cargo)
import numpy as np
import matplotlib.pyplot as plt
from scipy import interpolate
from matplotlib import cm
from scipy import ndimage

# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

# data input and output locations
dataPath = r'E:\PBAR\data4\BasicScanCargo'
dataOutputBaseDir = r'E:\PBAR\data4\BasicScansStandardWidth'

# define parameters for standard width images
numberTimeSlices = 2000
preCargoTimeSlices = 70


#################################
## DEFINE FILE NAMES

datStandardWidthList = []
datCargoStandardWidthList = []
markerStandardWidthList = []

datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')

for datasetIndex, datasetDescript in enumerate(datasetDescription['scanID']):    
    
#    if datasetIndex == 10:
#        break
    # open standard width files
    # Read Zspec
    filenameZspec = '%s-FDFC-All_SW.npy' %datasetDescription['scanID'][datasetIndex]
    fullFilenameZspec = os.path.join(dataPath, filenameZspec)
    print('Loading %s' %fullFilenameZspec)
    datStandardWidth = np.load(fullFilenameZspec)
    datStandardWidthList.append(datStandardWidth)
    
    # Read Cargo
    filenameCargo = 'PBAR-%s.cargoimageSW.npy' %datasetDescription['dataFile'][datasetIndex]
    fullFilenameCargo = os.path.join(dataPath, filenameCargo)
    print('Loading %s' %fullFilenameCargo)
    datCargoStandardWidth= np.load(fullFilenameCargo)
    datCargoStandardWidthList.append(datCargoStandardWidth)
    
    # Read in marker files
    filenameMarker = filenameCargo.replace('cargoimageSW', 'cargomarkerSW')
    fullFilenameMarker = fullFilenameCargo.replace('cargoimageSW', 'cargomarkerSW')
    # some don't have marker files
    if os.path.exists(fullFilenameMarker):
    #        markers = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
        markerStandardWidth = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
    else:
        markerStandardWidth = []
    markerStandardWidthList.append(markerStandardWidth)



# Plot of radiography vs zspec count for different lower bin threshold for zspec
# objective: determine what effect that bin threshold has on counts

plt.figure()
plt.grid()

datasetIndex = 10

datStandardWidth = datStandardWidthList[datasetIndex]
datCargoStandardWidth = datCargoStandardWidthList[datasetIndex]
markerStandardWidth = markerStandardWidthList[datasetIndex]

for minBin in [0, 6]:
    
    y = np.reshape(datStandardWidth[:,:,minBin:].sum(2), datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    x = np.reshape(datCargoStandardWidth, datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    
    plt.plot(x, y, '.k', color = plotColors[minBin], alpha = 0.1, label = 'min bin %d' %minBin)

plt.xlabel('Rad Count')
plt.ylabel('Zspec Count')
plt.legend(prop={'size':16})


# Plot of radiography vs zspec count for different lower bin threshold for zspec
# objective: determine what effect that bin threshold has on counts

plt.figure()
plt.grid()

datasetIndex = 8

datStandardWidth = datStandardWidthList[datasetIndex]
datCargoStandardWidth = datCargoStandardWidthList[datasetIndex]
markerStandardWidth = markerStandardWidthList[datasetIndex]

for minBin in [0, 1, 2, 3, 4, 5, 6]:
    
    y = np.reshape(datStandardWidth[:,:,minBin:].sum(2), datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    x = np.reshape(datCargoStandardWidth, datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    
    plt.plot(x, y, '.k', color = plotColors[minBin], alpha = 0.2, label = 'min bin %d' %minBin)

plt.xlabel('Rad Count')
plt.ylabel('Zspec Count')
plt.legend()

# plot of radiography vs zspec count
# plot all data points (not just ones in transmission range)
# objective - examine the region of linearity

# fit polynomial to the data points
pfitList = []
for datasetIndex in xrange(len(datStandardWidthList)):
    y = np.reshape(datStandardWidthList[datasetIndex][:,:,minBin:].sum(2), datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    x = np.reshape(datCargoStandardWidthList[datasetIndex], datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    cut = (x > 3.0e6) & (x < 3e7) & (y > 30) & (y < 350)
    pfitList.append(np.polyfit(x[cut], y[cut], 5))
# plot the cargo data points
plt.figure()
plt.grid()
minBin = 2
for datasetIndex in xrange(len(datStandardWidthList)):
    y = np.reshape(datStandardWidthList[datasetIndex][:,:,minBin:].sum(2), datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    x = np.reshape(datCargoStandardWidthList[datasetIndex], datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    
    plt.plot(x, y, '.k', color = plotColors[datasetIndex], alpha = 0.2, label = 'min bin %d, %d' %(minBin, datasetIndex))
    xx = np.linspace(0, 3e7, 50)
    yy = np.polyval(pfitList[datasetIndex], xx)
    plt.plot(xx, yy, color = plotColors[datasetIndex])

plt.xlabel('Rad Count')
plt.ylabel('Zspec Count')
plt.legend()
plt.axis( (0, 1e8, 0, 500))


figureSize = (16, 10)

plotDir = r'E:\PBAR\data4\BasicScanDiagnosticPlots\ZspecRadiographyLinearity'
for datasetIndex in xrange(len(datStandardWidthList)):
#    plt.figure()
    plt.figure(figsize = figureSize)
    plt.grid()
    minBin = 2
    y = np.reshape(datStandardWidthList[datasetIndex][:,:,minBin:].sum(2), datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    x = np.reshape(datCargoStandardWidthList[datasetIndex], datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    
    plt.plot(x, y, '.k', alpha = 0.2, label = 'min bin %d, set index %d, cargo %s, target %s' %(minBin, datasetIndex, datasetDescription['cargo'][datasetIndex], datasetDescription['targets'][datasetIndex]))
    xx = np.linspace(0, 3e7, 50)
    yy = np.polyval(pfitList[datasetIndex], xx)
    plt.plot(xx, yy, color = plotColors[datasetIndex])

    plt.xlabel('Rad Count')
    plt.ylabel('Zspec Count')
    plt.title('Config %d, Cargo %s, Target %s, %s_%s_ZspecRadiographyLinearity.png' \
        %(datasetDescription['config'][datasetIndex], datasetDescription['cargo'][datasetIndex], \
        datasetDescription['targets'][datasetIndex], datasetDescription['scanID'][datasetIndex], datasetDescription['dataFile'][datasetIndex]))
    plt.title
    plt.legend()
    plt.axis( (0, 1e8, 0, 500))
    plt.savefig(os.path.join(plotDir, '%s_%s_ZspecRadiographyLinearity.png' %(datasetDescription['scanID'][datasetIndex], datasetDescription['dataFile'][datasetIndex])))
    plt.close()
